# libs ----
library(sf)
library(geomander)
library(tidyverse)
library(redist)
library(ppmf)
library(grid)
library(patchwork)
library(wacolors)
library(doParallel)
library(foreach)
library(here)

# helper ----
source(here('R/00_custom_functions.R'))

# read in ----
ms <- readRDS(file = here('data/MS/ms.Rds'))

# set up maps ----
map_cen <- redist_map(ms, existing_plan = sld_up, pop_tol = 0.05, total_pop = pop)
map_v4  <- redist_map(ms, existing_plan = sld_up, pop_tol = 0.05, total_pop = v4_pop)
map_v12 <- redist_map(ms, existing_plan = sld_up, pop_tol = 0.05, total_pop = v12_pop)
map_v19 <- redist_map(ms, existing_plan = sld_up, pop_tol = 0.05, total_pop = v19_pop)


# with constraint:
cons_cen <- list(vra_old = list(strength = 100, tgt_vra_min = 0.55, tgt_vra_other = 0.2,
                                min_pop = ms$pop_black, pow_vra = 1.5))
cons_v4 <- list(vra_old = list(strength = 100, tgt_vra_min = 0.55, tgt_vra_other = 0.2,
                               min_pop = ms$v4_pop_black, pow_vra = 1.5))
cons_v12 <- list(vra_old = list(strength = 100, tgt_vra_min = 0.55, tgt_vra_other = 0.2,
                                min_pop = ms$v12_pop_black, pow_vra = 1.5))
cons_v19 <- list(vra_old = list(strength = 100, tgt_vra_min = 0.55, tgt_vra_other = 0.2,
                                min_pop = ms$v19_pop_black, pow_vra = 1.5))


# set up ----
nsim <- 5e4

# run
sims_cen<- redist_smc(map = map_cen, nsims = nsim, constraints = cons_cen,
                      truncate = TRUE)
sims_v4 <- redist_smc(map = map_v4, nsims = nsim, constraints = cons_v4,
                      truncate = TRUE)
sims_v12 <- redist_smc(map = map_v12, nsims = nsim, constraints = cons_v12,
                       truncate = TRUE)
sims_v19 <- redist_smc(map = map_v19, nsims = nsim, constraints = cons_v19,
                       truncate = TRUE)

# diversify
div_cen <- redist_mergesplit_parallel(map = map_cen, init_plan = get_plans_matrix(sims_cen)[,-1],
                                      nsims = 100, warmup = 99, constraints = cons_cen,
                                      ncores = 8, chains = nsim)
div_v4 <- redist_mergesplit_parallel(map = map_v4, init_plan = get_plans_matrix(sims_v4)[,-1],
                                     nsims = 100, warmup = 99, constraints = cons_v4,
                                     ncores = 8, chains = nsim)
div_v12 <- redist_mergesplit_parallel(map = map_v12, init_plan = get_plans_matrix(sims_v12)[,-1],
                                      nsims = 100, warmup = 99, constraints = cons_v12,
                                      ncores = 8, chains = nsim)
sims_v19 <- readRDS(here("data/MS/redist_plans_smc-trunc_v19.Rds"))
div_v19 <- redist_mergesplit_parallel(map = map_v19, init_plan = get_plans_matrix(sims_v19)[,-1],
                                      nsims = 100, warmup = 99, constraints = cons_v19,
                                      ncores = 8, chains = nsim)

saveRDS(div_cen, here('data/MS/sim/redist_plans_ms_cen.Rds'))
saveRDS(div_v4, here('data/MS/sim/redist_plans_ms_v4.Rds'))
saveRDS(div_v12, here('data/MS/sim/redist_plans_ms_v12.Rds'))
saveRDS(sims_v19, here('data/MS/sim/redist_plans_smc-trunc_v19.Rds'))
saveRDS(div_v19, here('data/MS/sim/redist_plans_ms_v19.Rds'))
